#include "ComputeProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "LocalShader.h"
#include "Utils/FileUtil.h"
#include "Game/GameState.h"
#include "Render/VulkanTool.h"


void ComputeProcess::OnInit()
{

	ENQUEUE_RENDER_COMMAND(BuildModel)([this](VkCommandBuffer commandBuffer)
		{
			mCommandBuffer = RHICommandBuffer::Builder()
				.Build();

			/*mLeftTexture = Picture2DTexture::Builder()
				.SetPath("VikingRoom.png")
				.SetMipLevels(1)
				.Build(this, "leftTexture");*/

			std::vector<uint8> leftBuffer;
			uint32 width = 0;
			uint32 height = 0;
			FileUtil::LoadTextureToArray("VikingRoom.png", VK_FORMAT_R8G8B8A8_UNORM, leftBuffer, width, height);

			mLeftTexture = CustomTexture::Builder()
				.SetWidth(width)
				.SetHeight(height)
				.SetFormat(VK_FORMAT_R8G8B8A8_UNORM)
				.SetUsage(VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_STORAGE_BIT)
				.SetMemoryProperty(VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT)
				.SetData(leftBuffer.data())
				.Build(this, "leftTexture");

			// Check if requested image format supports image storage operations
			VkFormatProperties formatProperties;
			vkGetPhysicalDeviceFormatProperties(GVulkanRHI->GetPhysicalDevice(), VK_FORMAT_R8G8B8A8_UNORM, &formatProperties);
			assert(formatProperties.optimalTilingFeatures & VK_FORMAT_FEATURE_STORAGE_IMAGE_BIT);

			mRightTexture = CustomTexture::Builder()
				.SetWidth(width)
				.SetHeight(height)
				.SetFormat(VK_FORMAT_R8G8B8A8_UNORM)
				.SetUsage(VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_STORAGE_BIT)
				.SetMemoryProperty(VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT)
				.AddQueueFamily(VK_QUEUE_GRAPHICS_BIT)
				.AddQueueFamily(VK_QUEUE_COMPUTE_BIT)
				.SetSharingMode(VK_SHARING_MODE_CONCURRENT)
				.Build(this, "rightTexture");

#if 0
			VkCommandBuffer transferCommandBuffer = GRenderState->BeginSingleCommand();

			VkImageMemoryBarrier imageMemoryBarrier{};
			imageMemoryBarrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
			imageMemoryBarrier.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT;
			imageMemoryBarrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
			imageMemoryBarrier.oldLayout = VK_IMAGE_LAYOUT_UNDEFINED;
			imageMemoryBarrier.newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
			imageMemoryBarrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
			imageMemoryBarrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
			imageMemoryBarrier.image = mRightTexture->GetImage();
			imageMemoryBarrier.subresourceRange = { VK_IMAGE_ASPECT_COLOR_BIT , 0, 1, 0, 1 };

			vkCmdPipelineBarrier(
				transferCommandBuffer,
				VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
				VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
				0,
				0, nullptr,
				0, nullptr,
				1, &imageMemoryBarrier
			);

			GRenderState->FlushSingleCommand(transferCommandBuffer);
#else
			TextureTransfer()
				.SetSrcStageMask(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_WRITE_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetNewLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.SetTexture(mRightTexture)
				.Transfer();
#endif

			// Default signaled, Compute process run first, Graphics wait
			mWaitGraphicsSemaphore = RHISemaphore::Builder()
				.SetWaitDstStage(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT)
				.SetCount(1)
				.SetSignaled(true)
				.Build();

			mSignalComputeSemaphore = RHISemaphore::Builder()
				.SetWaitDstStage(VK_PIPELINE_STAGE_VERTEX_INPUT_BIT)
				.SetCount(1)
				.Build();


			mEdgeDetectDescriptor = RHIDescriptor::Builder()
				.PushShader<EdgeDetectCompShader>()
				.SetDescriptorCount(0, 1, true)
				.Build();

			mEdgeDetectDescriptor->SetValue("inputImage", mLeftTexture, VK_IMAGE_LAYOUT_GENERAL);
			mEdgeDetectDescriptor->SetValue("resultImage", mRightTexture, VK_IMAGE_LAYOUT_GENERAL);
			mEdgeDetectDescriptorMask = mEdgeDetectDescriptor->AllocateDescriptor();

			mEdgeDetectPipeline = ComputePipeline::Builder()
				.SetDescriptor(mEdgeDetectDescriptor)
				.Build();




			mEmbossDescriptor = RHIDescriptor::Builder()
				.PushShader<EmbossCompShader>()
				.SetDescriptorCount(0, 1, true)
				.Build();

			mEmbossDescriptor->SetValue("inputImage", mLeftTexture, VK_IMAGE_LAYOUT_GENERAL);
			mEmbossDescriptor->SetValue("resultImage", mRightTexture, VK_IMAGE_LAYOUT_GENERAL);
			mEmbossDescriptorMask = mEmbossDescriptor->AllocateDescriptor();

			mEmbossPipeline = ComputePipeline::Builder()
				.SetDescriptor(mEmbossDescriptor)
				.Build();





			mSharpenDescriptor = RHIDescriptor::Builder()
				.PushShader<SharpenCompShader>()
				.SetDescriptorCount(0, 1, true)
				.Build();

			mSharpenDescriptor->SetValue("inputImage", mLeftTexture, VK_IMAGE_LAYOUT_GENERAL);
			mSharpenDescriptor->SetValue("resultImage", mRightTexture, VK_IMAGE_LAYOUT_GENERAL);
			mSharpenDescriptorMask = mSharpenDescriptor->AllocateDescriptor();

			mSharpenPipeline = ComputePipeline::Builder()
				.SetDescriptor(mSharpenDescriptor)
				.Build();



		});
}

void ComputeProcess::OnTick(float deltaTime)
{
	ENQUEUE_RENDER_COMMAND(RecordCommand)([this](VkCommandBuffer commandBuffer)
		{

			GRenderState->AddWaitSemaphore(mSignalComputeSemaphore);
			GRenderState->AddSignalSemaphore(mWaitGraphicsSemaphore);

			VkCommandBuffer computeCommandBuffer = mCommandBuffer->BeginCommandBuffer();

#if 0

			VkImageMemoryBarrier imageMemoryBarrier{};
			imageMemoryBarrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
			imageMemoryBarrier.srcAccessMask = VK_ACCESS_SHADER_READ_BIT;
			imageMemoryBarrier.dstAccessMask = VK_ACCESS_SHADER_WRITE_BIT;
			imageMemoryBarrier.oldLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
			imageMemoryBarrier.newLayout = VK_IMAGE_LAYOUT_GENERAL;
			imageMemoryBarrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
			imageMemoryBarrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
			imageMemoryBarrier.image = mLeftTexture->GetImage();
			imageMemoryBarrier.subresourceRange = { VK_IMAGE_ASPECT_COLOR_BIT , 0, 1, 0, 1 };

			vkCmdPipelineBarrier(
				computeCommandBuffer,
				VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
				VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
				0,
				0, nullptr,
				0, nullptr,
				1, &imageMemoryBarrier
			);

			imageMemoryBarrier.image = mRightTexture->GetImage();

			vkCmdPipelineBarrier(
				computeCommandBuffer,
				VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
				VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT,
				0,
				0, nullptr,
				0, nullptr,
				1, &imageMemoryBarrier
			);

#else

			TextureTransfer textureTransfer;
			textureTransfer
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_WRITE_BIT)
				.SetNewLayout(VK_IMAGE_LAYOUT_GENERAL)
				.SetTexture(mLeftTexture)
				.Transfer(computeCommandBuffer);

			textureTransfer
				.SetTexture(mRightTexture)
				.Transfer(computeCommandBuffer);

#endif

			// GRenderState->FlushSingleCommand(transferCommandBuffer);

			if (mPipelineIndex == 0)
			{
				vkCmdBindPipeline(computeCommandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, mEdgeDetectPipeline->GetPipeline());
				vkCmdBindDescriptorSets(computeCommandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, mEdgeDetectPipeline->GetPipelineLayout(), 0, 1, &mEdgeDetectDescriptorMask.GetDescriptorSet(), 0, 0);
			}
			else if (mPipelineIndex == 1)
			{
				vkCmdBindPipeline(computeCommandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, mEmbossPipeline->GetPipeline());
				vkCmdBindDescriptorSets(computeCommandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, mEmbossPipeline->GetPipelineLayout(), 0, 1, &mEmbossDescriptorMask.GetDescriptorSet(), 0, 0);
			}
			else if (mPipelineIndex == 2)
			{
				vkCmdBindPipeline(computeCommandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, mSharpenPipeline->GetPipeline());
				vkCmdBindDescriptorSets(computeCommandBuffer, VK_PIPELINE_BIND_POINT_COMPUTE, mSharpenPipeline->GetPipelineLayout(), 0, 1, &mSharpenDescriptorMask.GetDescriptorSet(), 0, 0);
			}

			vkCmdDispatch(computeCommandBuffer, mLeftTexture->GetWitdh() / 16, mLeftTexture->GetHeight() / 16, 1);

			mCommandBuffer->FlushCommandBuffer(computeCommandBuffer, { mWaitGraphicsSemaphore }, { mSignalComputeSemaphore });
			//mCommandBuffer->FlushCommandBuffer(computeCommandBuffer);
		});
}

void ComputeProcess::OnResize()
{

}
